const { loadModel, createCompletionGenerator } = require("gpt4all")
const { PassThrough } = require('stream')

const service = {
  model: null,
  chat: null,
  load: async (modelName) => {
    if(service.model){
      service.model.dispose()
      service.model = null
    }
    if(modelName.toLowerCase().endsWith('.gguf')){
      modelName = modelName.substr(0, modelName.length - 5)
    }
    service.model = await loadModel(modelName, { modelPath: __dirname + "/models", verbose: false, device: 'gpu', modelConfigFile: "./models.json" })
  },
  generate: async(modelName, messages) => {
    if(!service.model){
      await service.load(modelName)
    }
    const stream = new PassThrough()
    const asyncGen = createCompletionGenerator(service.model, messages)
    const completion = async() => {
      while (true) {
        const { value, done } = await asyncGen.next()
        if (done) { break; }
        const data = { choices: [{ delta: { content: value } }] }
        stream.write(`data: ${JSON.stringify(data)}\n`)
      }
    }
    completion().then(()=>{
      stream.write(`data: [DONE]\n`)
      stream.end()
    })
    return stream
  },
  complete: async(sessionId, modelName, messages) => {
    if(!service.model){
      await service.load(modelName)
    }
    if(!service.chat || service.chat.id != sessionId){
      const system = messages[0]
      const historys = messages.slice(1).slice(0, -1).slice(-20)
      if(historys.length > 0){
        service.chat = await service.model.createChatSession({ systemPrompt: system.content, messages: historys })
      }else if(historys.length == 0){
        service.chat = await service.model.createChatSession({ systemPrompt: system.content })
      }else{
        service.chat = await service.model.createChatSession()
      }
      service.chat.id = sessionId
    }
    let input = messages[messages.length - 1]
    const stream = new PassThrough()
    const asyncGen = createCompletionGenerator(service.chat, input.content)
    const completion = async() => {
      while (true) {
        const { value, done } = await asyncGen.next()
        if (done) { break; }
        const data = { choices: [{ delta: { content: value } }] }
        stream.write(`data: ${JSON.stringify(data)}\n`)
      }
    }
    completion().then(()=>{
      stream.write(`data: [DONE]\n`)
      stream.end()
    })
    return stream
  },
}

module.exports = service